#include "fetchdialog.h"

#include <QDialogButtonBox>
#include <QFormLayout>
#include <QInputDialog>
#include <QLabel>

#include "action.h"
#include "fetcher.h"
#include "misc.h"

using namespace BinaryNinja;

static void AddListItem(QListWidget* list, const QString& value)
{
	if (value.trimmed().isEmpty())
		return;
	// Avoid duplicates
	for (int i = 0; i < list->count(); ++i)
		if (list->item(i)->text().compare(value, Qt::CaseInsensitive) == 0)
			return;
	list->addItem(value.trimmed());
}

WarpFetchDialog::WarpFetchDialog(BinaryViewRef bv, std::shared_ptr<WarpFetcher> fetcher, QWidget* parent) :
	QDialog(parent), m_fetchProcessor(std::move(fetcher)), m_bv(std::move(bv))
{
	setWindowTitle("Fetch WARP Functions");

	auto form = new QFormLayout();
	m_containerCombo = new QComboBox(this);
	populateContainers();
	m_containerCombo->addItem("All Containers");  // index 0 for "all"
	for (const auto& c : m_containers)
		m_containerCombo->addItem(QString::fromStdString(c->GetName()));

	// Tags editor
	m_tagsList = new QListWidget(this);
	m_addTagBtn = new QPushButton(this);
	m_addTagBtn->setText("+");
	m_addTagBtn->setToolTip("Add tag");
	m_removeTagBtn = new QPushButton(this);
	m_removeTagBtn->setText("-");
	m_removeTagBtn->setToolTip("Remove selected tag(s)");
	m_resetTagBtn = new QPushButton(this);
	m_resetTagBtn->setText("Reset");
	m_resetTagBtn->setToolTip("Reset tags to: official, trusted");
	auto tagBtnRow = new QHBoxLayout();
	tagBtnRow->addWidget(m_addTagBtn);
	tagBtnRow->addWidget(m_removeTagBtn);
	tagBtnRow->addWidget(m_resetTagBtn);
	tagBtnRow->addStretch();
	auto tagCol = new QVBoxLayout();
	tagCol->addWidget(m_tagsList);
	tagCol->addLayout(tagBtnRow);
	auto tagWrapper = new QWidget(this);
	tagWrapper->setLayout(tagCol);

	// Make tags list compact with a fixed maximum height and no vertical expansion
	m_tagsList->setSizeAdjustPolicy(QAbstractScrollArea::AdjustToContents);
	m_tagsList->setVerticalScrollBarPolicy(Qt::ScrollBarAsNeeded);
	m_tagsList->setSizePolicy(QSizePolicy::Expanding, QSizePolicy::Fixed);
	m_tagsList->setMaximumHeight(120);
	m_tagsList->setToolTip("A source must have atleast ONE of these tags to be considered");

	// Defaults from processor tags
	for (const auto& t : GetAllowedTagsFromView(m_bv))
		AddListItem(m_tagsList, QString::fromStdString(t));

	// Batch size and matcher checkbox
	m_batchSize = new QSpinBox(this);
	m_batchSize->setRange(10, 1000);
	m_batchSize->setValue(GetBatchSizeFromView(m_bv));
	m_batchSize->setToolTip("Number of functions to fetch in each batch");

	m_rerunMatcher = new QCheckBox("Re-run matcher after fetch", this);
	m_rerunMatcher->setChecked(true);

	m_clearProcessed = new QCheckBox("Refetch all functions", this);
	m_clearProcessed->setToolTip(
		"Clears the processed cache before fetching again, this will refetch all functions in the view");
	m_clearProcessed->setChecked(false);

	form->addRow(new QLabel("Container: "), m_containerCombo);
	form->addRow(new QLabel("Allowed Tags: "), tagWrapper);
	form->addRow(new QLabel("Batch Size: "), m_batchSize);
	form->addRow(m_rerunMatcher);
	form->addRow(m_clearProcessed);

	auto buttons = new QDialogButtonBox(QDialogButtonBox::Ok | QDialogButtonBox::Cancel, this);
	connect(buttons, &QDialogButtonBox::accepted, this, &WarpFetchDialog::onAccept);
	connect(buttons, &QDialogButtonBox::rejected, this, &WarpFetchDialog::onReject);

	auto root = new QVBoxLayout(this);
	root->addLayout(form);
	root->addWidget(buttons);
	setLayout(root);

	// Wire buttons
	connect(m_addTagBtn, &QPushButton::clicked, this, &WarpFetchDialog::onAddTag);
	connect(m_removeTagBtn, &QPushButton::clicked, this, &WarpFetchDialog::onRemoveTag);
	connect(m_resetTagBtn, &QPushButton::clicked, this, &WarpFetchDialog::onResetTags);
}

void WarpFetchDialog::populateContainers()
{
	m_containers = Warp::Container::All();
}

void WarpFetchDialog::onAddTag()
{
	bool ok = false;
	const auto text = QInputDialog::getText(this, "Add Tag", "Tag:", QLineEdit::Normal, {}, &ok);
	if (ok)
		AddListItem(m_tagsList, text);
}

void WarpFetchDialog::onRemoveTag()
{
	for (auto* item : m_tagsList->selectedItems())
		delete item;
}

void WarpFetchDialog::onResetTags()
{
	m_tagsList->clear();
	AddListItem(m_tagsList, "official");
	AddListItem(m_tagsList, "trusted");
}

std::vector<Warp::SourceTag> WarpFetchDialog::collectTags() const
{
	std::vector<Warp::SourceTag> out;
	out.reserve(m_tagsList->count());
	for (int i = 0; i < m_tagsList->count(); ++i)
		out.emplace_back(m_tagsList->item(i)->text().trimmed().toStdString());
	return out;
}

void WarpFetchDialog::onAccept()
{
	const int idx = m_containerCombo->currentIndex();
	std::optional<size_t> containerIndex;
	if (idx > 0)  // 0 == All Containers
		containerIndex = static_cast<size_t>(idx - 1);

	const auto batch = static_cast<size_t>(m_batchSize->value());
	const bool rerun = m_rerunMatcher->isChecked();

	const auto tags = collectTags();
	// Persist tags to the view settings.
	SetTagsToView(m_bv, tags);

	if (m_clearProcessed->isChecked())
		m_fetchProcessor->ClearProcessed();

	// Execute the network fetch in batches
	runBatchedFetch(containerIndex, tags, batch, rerun);

	accept();
}

void WarpFetchDialog::onReject()
{
	const auto tags = collectTags();
	// Persist tags to the view settings.
	SetTagsToView(m_bv, tags);
	reject();
}

void WarpFetchDialog::runBatchedFetch(const std::optional<size_t>& containerIndex,
	const std::vector<Warp::SourceTag>& allowedTags, size_t batchSize, bool rerunMatcher)
{
	if (!m_bv)
		return;
	// Collect functions in the view and enqueue them to the shared fetcher
	std::vector<Ref<Function>> funcs = m_bv->GetAnalysisFunctionList();
	if (funcs.empty())
		return;
	const size_t totalFuncs = funcs.size();
	const size_t totalBatches = (totalFuncs + batchSize - 1) / batchSize;

	// Create a background task to show progress in the UI
	Ref<BackgroundTask> task =
		new BackgroundTask("Fetching WARP functions (0 / " + std::to_string(totalBatches) + ")", false);

	auto fetcher = m_fetchProcessor;
	auto bv = m_bv;

	// TODO: Too many captures in this thing lol.
	WorkerInteractiveEnqueue(
		[fetcher, bv, funcs = std::move(funcs), batchSize, rerunMatcher, task, allowedTags]() mutable {
			size_t processed = 0;
			size_t batchIndex = 0;

			while (processed < funcs.size())
			{
				const size_t remaining = funcs.size() - processed;
				const size_t thisBatchCount = std::min(batchSize, remaining);

				for (size_t i = 0; i < thisBatchCount; ++i)
					fetcher->AddPendingFunction(funcs[processed + i]);

				fetcher->FetchPendingFunctions(allowedTags);

				++batchIndex;
				processed += thisBatchCount;

				task->SetProgressText("Fetching WARP functions (" + std::to_string(batchIndex) + " / "
					+ std::to_string((funcs.size() + batchSize - 1) / batchSize) + ")");
			}

			task->Finish();
			// TODO: Print how long it took?
			Logger("WARP Fetcher").LogInfo("Finished fetching WARP functions...");

			if (rerunMatcher && bv)
				Warp::RunMatcher(*bv);
		});
}

void RegisterWarpFetchFunctionsCommand()
{
	// Register a UI action and bind it globally. Add it to the Tools menu.
	const QString actionName = "WARP\\Fetch";
	if (!UIAction::isActionRegistered(actionName))
		UIAction::registerAction(actionName);

	UIActionHandler::globalActions()->bindAction(actionName,
		UIAction(
			[](const UIActionContext& context) {
				if (const BinaryViewRef bv = context.binaryView; bv)
				{
					WarpFetchDialog dlg(bv, WarpFetcher::Global(), nullptr);
					dlg.exec();
				}
			},
			[](const UIActionContext& context) { return context.binaryView != nullptr; }));

	Menu::mainMenu("Plugins")->addAction(actionName, "Plugins");
}
